package cn.kimmking.kksharding.mybatis;

import cn.kimmking.kksharding.engine.ShardingContext;
import cn.kimmking.kksharding.engine.ShardingResult;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.springframework.objenesis.instantiator.util.UnsafeUtils;
import sun.misc.Unsafe;

import java.lang.reflect.Field;

/**
 * intercept sql.
 *
 * @Author : kimmking(kimmking@apache.org)
 * @create 2024/7/22 下午8:26
 */

@Intercepts(
        @org.apache.ibatis.plugin.Signature(
                type = StatementHandler.class,
                method = "prepare",
                args = {java.sql.Connection.class, Integer.class}
        )
)
public class SqlStatementInterceptor implements Interceptor {
        @Override
        public Object intercept(Invocation invocation) throws Throwable {
                ShardingResult result = ShardingContext.get();
                if(result != null) {
                        StatementHandler handler = (StatementHandler) invocation.getTarget();
                        BoundSql boundSql = handler.getBoundSql();
                        String sql = boundSql.getSql();
                        System.out.println(" ===> SqlStatementInterceptor: " + sql);
                        String targetSqlStatement = result.getTargetSqlStatement();
                        if(!sql.equalsIgnoreCase(targetSqlStatement)) {
                                replaceSql(boundSql, targetSqlStatement);
                        }
                }
                return invocation.proceed();
        }

        private static void replaceSql(BoundSql boundSql, String sql) throws NoSuchFieldException {
                Field field = boundSql.getClass().getDeclaredField("sql");
                Unsafe unsafe = UnsafeUtils.getUnsafe();
                long fieldOffset = unsafe.objectFieldOffset(field);
                unsafe.putObject(boundSql, fieldOffset, sql);
        }
}
